import os
import matplotlib.pyplot as plt
import json

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 获取当前脚本的绝对路径，用于排除自身
CURRENT_SCRIPT_PATH = os.path.abspath(__file__)

def count_code_lines_in_notebook(filepath):
    """统计Jupyter Notebook文件中的代码行数"""
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            notebook = json.load(f)
            
        code_lines = 0
        # 遍历notebook中的所有单元格
        for cell in notebook.get('cells', []):
            # 只统计代码单元格
            if cell.get('cell_type') == 'code':
                # 统计非空行
                for line in cell.get('source', []):
                    if line.strip():  # 非空行
                        code_lines += 1
        return code_lines
    except Exception as e:
        print(f"无法读取Notebook文件 {filepath}: {e}")
        return 0

def count_files_and_lines(root_dir, exclude_file):
    """统计文件夹中各类文件的数量和Python代码行数"""
    # 初始化计数器
    counts = {
        'python_files': 0,
        'notebook_files': 0,
        'gif_files': 0,
        'image_files': 0,
        'python_lines': 0
    }
    
    # 支持的图片格式
    image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.svg'}
    
    # 遍历目录
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for filename in filenames:
            filepath = os.path.join(dirpath, filename)
            ext = os.path.splitext(filename)[1].lower()
            
            # 跳过指定的文件（通常是本脚本文件）
            if os.path.abspath(filepath) == exclude_file:
                continue
            
            # 统计Python文件
            if ext == '.py':
                counts['python_files'] += 1
                # 统计代码行数
                try:
                    with open(filepath, 'r', encoding='utf-8') as f:
                        lines = f.readlines()
                        # 只统计非空行和非注释行
                        code_lines = [line for line in lines if line.strip() and not line.strip().startswith('#')]
                        counts['python_lines'] += len(code_lines)
                except Exception as e:
                    print(f"无法读取文件 {filepath}: {e}")
            
            # 统计Jupyter Notebook文件
            elif ext == '.ipynb':
                counts['notebook_files'] += 1
                # 统计Notebook中的代码行数
                notebook_code_lines = count_code_lines_in_notebook(filepath)
                counts['python_lines'] += notebook_code_lines
            
            # 统计GIF文件
            elif ext == '.gif':
                counts['gif_files'] += 1
            
            # 统计其他图片文件
            elif ext in image_extensions:
                counts['image_files'] += 1
    
    return counts

def create_bar_chart(counts, output_path):
    """创建柱状图并保存"""
    # 准备数据（不包含代码行数）
    labels = ['Python文件', 'Notebook文件', 'GIF动画', '图片文件']
    values = [counts['python_files'], counts['notebook_files'], counts['gif_files'], counts['image_files']]
    
    # 创建柱状图
    plt.figure(figsize=(8, 4))
    bars = plt.bar(labels, values, color=['blue', 'purple', 'green', 'orange'])
    
    # 添加数值标签
    for bar, value in zip(bars, values):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height(), str(value), 
                ha='center', va='bottom')
    
    # 在标题中显示代码总行数
    plt.title(f'文件统计结果 (代码总行数: {counts["python_lines"]})')
    plt.ylabel('数量')
    plt.xticks(rotation=45)
    plt.tight_layout()
    
    # 保存图片
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"柱状图已保存到: {output_path}")

def main():
    # 设置根目录
    root_dir = r'd:\yun\gitee\MathPhys'
    
    # 统计文件，排除本脚本文件
    counts = count_files_and_lines(root_dir, CURRENT_SCRIPT_PATH)
    
    # 打印统计结果
    print("文件统计结果:")
    print(f"Python文件数量: {counts['python_files']}")
    print(f"Notebook文件数量: {counts['notebook_files']}")
    print(f"GIF动画数量: {counts['gif_files']}")
    print(f"图片文件数量: {counts['image_files']}")
    print(f"代码总行数: {counts['python_lines']}")
    
    # 创建并保存柱状图
    output_path = os.path.join(root_dir, '文件统计柱状图.png')
    create_bar_chart(counts, output_path)

if __name__ == "__main__":
    main()